# -*- coding: utf-8 -*-
"""
提取天津2020年ERA5日降水量数据并导出到Excel

此脚本根据天津的经纬度点位，从ERA5.P.daily.2020.nc文件中提取2020年的降水数据，
并为每个经纬度点位生成对应的Excel工作表，工作表名称为经度_纬度。
"""

import os
import pandas as pd
import xarray as xr
from parse_era5_data import create_parser, DailyParser

# 文件路径配置
GRID_POINTS_FILE = "../province_grid_points/天津_grid_points.csv"
ERA5_DATA_FILE = "../Daily_ERA5(1960-2024)/ERA5.P.daily.2020.nc"
OUTPUT_EXCEL = "天津_降水量_2020.xlsx"

def load_grid_points(file_path):
    """
    加载经纬度点位数据
    
    参数:
        file_path (str): CSV文件路径
        
    返回:
        pandas.DataFrame: 包含经纬度点位的数据框
    """
    try:
        df = pd.read_csv(file_path)
        print(f"成功加载点位数据，共{len(df)}个点位")
        return df
    except Exception as e:
        print(f"加载点位数据时出错: {e}")
        return None

def extract_precipitation_data(era5_file, grid_points, output_excel):
    """
    提取降水数据并导出到Excel
    
    参数:
        era5_file (str): ERA5数据文件路径
        grid_points (pandas.DataFrame): 经纬度点位数据
        output_excel (str): 输出Excel文件路径
        
    返回:
        bool: 成功返回True，失败返回False
    """
    try:
        # 创建ERA5解析器
        parser = create_parser(era5_file)
        if parser is None:
            print("无法创建解析器，程序退出")
            return False
            
        # 确保是日降水量数据解析器
        if not isinstance(parser, DailyParser):
            print("数据文件不是日降水量格式，程序退出")
            parser.close()
            return False
            
        # 显示数据集基本信息
        parser.display_info()
        
        # 获取主要变量名称（降水量）
        main_var_name = list(parser.dataset.data_vars)[0]
        print(f"提取变量: {main_var_name}")
        
        # 获取空间维度名称
        lat_dim, lon_dim = parser.get_spatial_dimensions()
        
        # 创建ExcelWriter对象
        with pd.ExcelWriter(output_excel, engine='openpyxl') as writer:
            # 遍历每个经纬度点位
            for index, row in grid_points.iterrows():
                lon = row['经度']
                lat = row['纬度']
                
                print(f"处理点位 {index+1}/{len(grid_points)}: 经度={lon}, 纬度={lat}")
                
                # 提取该点位的数据
                data_array = parser.extract_variable(main_var_name)
                
                # 选择最接近的经纬度点
                data_point = data_array.sel({lat_dim: lat, lon_dim: lon}, method='nearest')
                
                # 获取实际选择的经纬度值
                actual_lat = float(data_point[lat_dim].values)
                actual_lon = float(data_point[lon_dim].values)
                print(f"实际点位: 经度={actual_lon}, 纬度={actual_lat}")
                
                # 转换为DataFrame
                df = data_point.to_dataframe().reset_index()
                
                # 添加日期列
                df['日期'] = df[parser.time_dim].apply(lambda x: parser.format_date(x))
                
                # 重新排列列顺序
                df = df[['日期', main_var_name]]
                
                # 重命名列
                df.rename(columns={main_var_name: '降水量'}, inplace=True)
                
                # 创建工作表名称（经度_纬度）
                sheet_name = f"{lon}_{lat}"
                
                # 导出到Excel工作表
                df.to_excel(writer, sheet_name=sheet_name, index=False)
                print(f"已导出到工作表: {sheet_name}")
            
            print(f"\n所有数据已成功导出到: {output_excel}")
        
        # 关闭数据集
        parser.close()
        return True
        
    except Exception as e:
        print(f"提取数据时出错: {e}")
        return False

def main():
    """
    主函数
    """
    print("开始提取天津2020年ERA5日降水量数据...")
    print("=" * 50)
    
    # 获取当前工作目录
    work_dir = os.getcwd()
    print(f"当前工作目录: {work_dir}")
    
    # 构建完整文件路径
    grid_points_file = os.path.join(work_dir, GRID_POINTS_FILE)
    era5_file = os.path.join(work_dir, ERA5_DATA_FILE)
    output_excel = os.path.join(work_dir, OUTPUT_EXCEL)
    
    # 加载经纬度点位数据
    grid_points = load_grid_points(grid_points_file)
    if grid_points is None:
        return
    
    # 提取降水数据并导出到Excel
    success = extract_precipitation_data(era5_file, grid_points, output_excel)
    
    if success:
        print("\n程序执行完毕，数据已成功导出")
    else:
        print("\n程序执行失败")

if __name__ == "__main__":
    main()